// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/io/buffered.h>

#include <algorithm>
#include <cstring>
#include <memory>
#include <mutex>
#include <string_view>
#include <utility>

#include <nebula/core/buffer.h>
#include <nebula/io/util_internal.h>
#include <nebula/core/memory_pool.h>
#include <turbo/utility/status.h>
#include <turbo/log/logging.h>

namespace nebula::io {

    // ----------------------------------------------------------------------
    // BufferedOutputStream implementation

    class BufferedBase {
    public:
        explicit BufferedBase(MemoryPool *pool)
                : pool_(pool),
                  is_open_(true),
                  buffer_data_(nullptr),
                  buffer_pos_(0),
                  buffer_size_(0),
                  raw_pos_(-1) {}

        bool closed() const {
            std::lock_guard<std::mutex> guard(lock_);
            return !is_open_;
        }

        turbo::Status ResetBuffer() {
            if (!buffer_) {
                // On first invocation, or if the buffer has been released, we allocate a
                // new buffer
                TURBO_MOVE_OR_RAISE(buffer_, allocate_resizable_buffer(buffer_size_, pool_));
            } else if (buffer_->size() != buffer_size_) {
                TURBO_RETURN_NOT_OK(buffer_->resize(buffer_size_));
            }
            buffer_data_ = buffer_->mutable_data();
            return turbo::OkStatus();
        }

        turbo::Status ResizeBuffer(int64_t new_buffer_size) {
            buffer_size_ = new_buffer_size;
            return ResetBuffer();
        }

        void AppendToBuffer(const void *data, int64_t nbytes) {
                    DKCHECK_LE(buffer_pos_ + nbytes, buffer_size_);
            std::memcpy(buffer_data_ + buffer_pos_, data, nbytes);
            buffer_pos_ += nbytes;
        }

        int64_t buffer_size() const { return buffer_size_; }

        int64_t buffer_pos() const { return buffer_pos_; }

    protected:
        MemoryPool *pool_;
        bool is_open_;

        std::shared_ptr<ResizableBuffer> buffer_;
        uint8_t *buffer_data_;
        int64_t buffer_pos_;
        int64_t buffer_size_;

        mutable int64_t raw_pos_;
        mutable std::mutex lock_;
    };

    class BufferedOutputStream::Impl : public BufferedBase {
    public:
        explicit Impl(std::shared_ptr<OutputStream> raw, MemoryPool *pool)
                : BufferedBase(pool), raw_(std::move(raw)) {}

        turbo::Status close() {
            std::lock_guard<std::mutex> guard(lock_);
            if (is_open_) {
                turbo::Status st = FlushUnlocked();
                is_open_ = false;
                TURBO_RETURN_NOT_OK(raw_->close());
                return st;
            }
            return turbo::OkStatus();
        }

        turbo::Status abort() {
            std::lock_guard<std::mutex> guard(lock_);
            if (is_open_) {
                is_open_ = false;
                return raw_->abort();
            }
            return turbo::OkStatus();
        }

        turbo::Result<int64_t> tell() const {
            std::lock_guard<std::mutex> guard(lock_);
            if (raw_pos_ == -1) {
                TURBO_MOVE_OR_RAISE(raw_pos_, raw_->tell());
                        DKCHECK_GE(raw_pos_, 0);
            }
            return raw_pos_ + buffer_pos_;
        }

        turbo::Status write(const void *data, int64_t nbytes) { return DoWrite(data, nbytes); }

        turbo::Status write(const std::shared_ptr<Buffer> &buffer) {
            return DoWrite(buffer->data(), buffer->size(), buffer);
        }

        turbo::Status DoWrite(const void *data, int64_t nbytes,
                       const std::shared_ptr<Buffer> &buffer = nullptr) {
            std::lock_guard<std::mutex> guard(lock_);
            if (nbytes < 0) {
                return turbo::invalid_argument_error("write count should be >= 0");
            }
            if (nbytes == 0) {
                return turbo::OkStatus();
            }
            if (nbytes + buffer_pos_ >= buffer_size_) {
                TURBO_RETURN_NOT_OK(FlushUnlocked());
                        DKCHECK_EQ(buffer_pos_, 0);
                if (nbytes >= buffer_size_) {
                    // Invalidate cached raw pos
                    raw_pos_ = -1;
                    // Direct write
                    if (buffer) {
                        return raw_->write(buffer);
                    } else {
                        return raw_->write(data, nbytes);
                    }
                }
            }
            AppendToBuffer(data, nbytes);
            return turbo::OkStatus();
        }

        turbo::Status FlushUnlocked() {
            if (buffer_pos_ > 0) {
                // Invalidate cached raw pos
                raw_pos_ = -1;
                TURBO_RETURN_NOT_OK(raw_->write(buffer_data_, buffer_pos_));
                buffer_pos_ = 0;
            }
            return turbo::OkStatus();
        }

        turbo::Status flush() {
            std::lock_guard<std::mutex> guard(lock_);
            return FlushUnlocked();
        }

        turbo::Result<std::shared_ptr<OutputStream>> detach() {
            std::lock_guard<std::mutex> guard(lock_);
            TURBO_RETURN_NOT_OK(FlushUnlocked());
            is_open_ = false;
            return std::move(raw_);
        }

        turbo::Status set_buffer_size(int64_t new_buffer_size) {
            std::lock_guard<std::mutex> guard(lock_);
            if (new_buffer_size <= 0) {
                return turbo::invalid_argument_error("Buffer size should be positive");
            }
            if (buffer_pos_ >= new_buffer_size) {
                // If the buffer is shrinking, first flush to the raw OutputStream
                TURBO_RETURN_NOT_OK(FlushUnlocked());
            }
            return ResizeBuffer(new_buffer_size);
        }

        std::shared_ptr<OutputStream> raw() const { return raw_; }

    private:
        std::shared_ptr<OutputStream> raw_;
    };

    BufferedOutputStream::BufferedOutputStream(std::shared_ptr<OutputStream> raw,
                                               MemoryPool *pool) {
        impl_.reset(new Impl(std::move(raw), pool));
    }

    turbo::Result<std::shared_ptr<BufferedOutputStream>> BufferedOutputStream::create(
            int64_t buffer_size, MemoryPool *pool, std::shared_ptr<OutputStream> raw) {
        auto result = std::shared_ptr<BufferedOutputStream>(
                new BufferedOutputStream(std::move(raw), pool));
        TURBO_RETURN_NOT_OK(result->set_buffer_size(buffer_size));
        return result;
    }

    BufferedOutputStream::~BufferedOutputStream() { internal::CloseFromDestructor(this); }

    turbo::Status BufferedOutputStream::set_buffer_size(int64_t new_buffer_size) {
        return impl_->set_buffer_size(new_buffer_size);
    }

    int64_t BufferedOutputStream::buffer_size() const { return impl_->buffer_size(); }

    int64_t BufferedOutputStream::bytes_buffered() const { return impl_->buffer_pos(); }

    turbo::Result<std::shared_ptr<OutputStream>> BufferedOutputStream::detach() {
        return impl_->detach();
    }

    turbo::Status BufferedOutputStream::close() { return impl_->close(); }

    turbo::Status BufferedOutputStream::abort() { return impl_->abort(); }

    bool BufferedOutputStream::closed() const { return impl_->closed(); }

    turbo::Result<int64_t> BufferedOutputStream::tell() const { return impl_->tell(); }

    turbo::Status BufferedOutputStream::write(const void *data, int64_t nbytes) {
        return impl_->write(data, nbytes);
    }

    turbo::Status BufferedOutputStream::write(const std::shared_ptr<Buffer> &data) {
        return impl_->write(data);
    }

    turbo::Status BufferedOutputStream::flush() { return impl_->flush(); }

    std::shared_ptr<OutputStream> BufferedOutputStream::raw() const { return impl_->raw(); }

    // ----------------------------------------------------------------------
    // BufferedInputStream implementation

    class BufferedInputStream::Impl : public BufferedBase {
    public:
        Impl(std::shared_ptr<InputStream> raw, MemoryPool *pool, int64_t raw_total_bytes_bound)
                : BufferedBase(pool),
                  raw_(std::move(raw)),
                  raw_read_total_(0),
                  raw_read_bound_(raw_total_bytes_bound),
                  bytes_buffered_(0) {}

        turbo::Status close() {
            if (is_open_) {
                is_open_ = false;
                return raw_->close();
            }
            return turbo::OkStatus();
        }

        turbo::Status abort() {
            if (is_open_) {
                is_open_ = false;
                return raw_->abort();
            }
            return turbo::OkStatus();
        }

        turbo::Result<int64_t> tell() const {
            if (raw_pos_ == -1) {
                TURBO_MOVE_OR_RAISE(raw_pos_, raw_->tell());
                        DKCHECK_GE(raw_pos_, 0);
            }
            // Shift by bytes_buffered to return semantic stream position
            return raw_pos_ - bytes_buffered_;
        }

        // resize internal read buffer. Note that the internal buffer-size
        // should be not larger than the raw_read_bound_.
        turbo::Status set_buffer_size(int64_t new_buffer_size) {
            if (new_buffer_size <= 0) {
                return turbo::invalid_argument_error("Buffer size should be positive");
            }
            if ((buffer_pos_ + bytes_buffered_) >= new_buffer_size) {
                return turbo::invalid_argument_error("Cannot shrink read buffer if buffered data remains");
            }
            if (raw_read_bound_ >= 0) {
                // No need to reserve space for more than the total remaining number of bytes.
                new_buffer_size = std::min(new_buffer_size,
                                           bytes_buffered_ + (raw_read_bound_ - raw_read_total_));
            }
            return ResizeBuffer(new_buffer_size);
        }

        turbo::Result<std::string_view> peek(int64_t nbytes) {
            if (raw_read_bound_ >= 0) {
                // Do not try to peek more than the total remaining number of bytes.
                nbytes = std::min(nbytes, bytes_buffered_ + (raw_read_bound_ - raw_read_total_));
            }

            if (bytes_buffered_ == 0 && nbytes < buffer_size_) {
                // Pre-buffer for small reads
                TURBO_RETURN_NOT_OK(BufferIfNeeded());
            }

            // Increase the buffer size if needed.
            if (nbytes > buffer_->size() - buffer_pos_) {
                TURBO_RETURN_NOT_OK(set_buffer_size(nbytes + buffer_pos_));
                        DKCHECK(buffer_->size() - buffer_pos_ >= nbytes);
            }
            // read more data when buffer has insufficient left
            if (nbytes > bytes_buffered_) {
                int64_t additional_bytes_to_read = nbytes - bytes_buffered_;
                if (raw_read_bound_ >= 0) {
                    additional_bytes_to_read =
                            std::min(additional_bytes_to_read, raw_read_bound_ - raw_read_total_);
                }
                TURBO_MOVE_OR_RAISE(
                        int64_t bytes_read,
                        raw_->read(additional_bytes_to_read,
                                   buffer_->mutable_data() + buffer_pos_ + bytes_buffered_));
                bytes_buffered_ += bytes_read;
                raw_read_total_ += bytes_read;
                nbytes = bytes_buffered_;
            }
                    DKCHECK(nbytes <= bytes_buffered_);  // Enough bytes available
            return std::string_view(reinterpret_cast<const char *>(buffer_data_ + buffer_pos_),
                                    static_cast<size_t>(nbytes));
        }

        int64_t bytes_buffered() const { return bytes_buffered_; }

        int64_t buffer_size() const { return buffer_size_; }

        std::shared_ptr<InputStream> detach() {
            is_open_ = false;
            return std::move(raw_);
        }

        void RewindBuffer() {
            // Invalidate buffered data, as with a Seek or large read
            buffer_pos_ = bytes_buffered_ = 0;
        }

        turbo::Status DoBuffer() {
            // Fill buffer
            if (!buffer_) {
                TURBO_RETURN_NOT_OK(ResetBuffer());
            }

            int64_t bytes_to_buffer = buffer_size_;
            if (raw_read_bound_ >= 0) {
                bytes_to_buffer = std::min(buffer_size_, raw_read_bound_ - raw_read_total_);
            }
            TURBO_MOVE_OR_RAISE(bytes_buffered_, raw_->read(bytes_to_buffer, buffer_data_));
            buffer_pos_ = 0;
            raw_read_total_ += bytes_buffered_;

            // Do not make assumptions about the raw stream position
            raw_pos_ = -1;
            return turbo::OkStatus();
        }

        turbo::Status BufferIfNeeded() {
            if (bytes_buffered_ == 0) {
                return DoBuffer();
            }
            return turbo::OkStatus();
        }

        void ConsumeBuffer(int64_t nbytes) {
            buffer_pos_ += nbytes;
            bytes_buffered_ -= nbytes;
        }

        turbo::Result<int64_t> read(int64_t nbytes, void *out) {
            if (TURBO_UNLIKELY(nbytes < 0)) {
                return turbo::invalid_argument_error("Bytes to read must be positive. Received:", nbytes);
            }

            // 1. First consume pre-buffered data.
            int64_t pre_buffer_copy_bytes = std::min(nbytes, bytes_buffered_);
            if (pre_buffer_copy_bytes > 0) {
                memcpy(out, buffer_data_ + buffer_pos_, pre_buffer_copy_bytes);
                ConsumeBuffer(pre_buffer_copy_bytes);
            }
            int64_t remaining_bytes = nbytes - pre_buffer_copy_bytes;
            if (raw_read_bound_ >= 0) {
                remaining_bytes = std::min(remaining_bytes, raw_read_bound_ - raw_read_total_);
            }
            if (remaining_bytes == 0) {
                return pre_buffer_copy_bytes;
            }
                    DKCHECK_EQ(0, bytes_buffered_);

            // 2. read from storage.
            if (remaining_bytes >= buffer_size_) {
                // 2.1. If read is larger than buffer size, read directly from storage.
                TURBO_MOVE_OR_RAISE(int64_t bytes_read,
                                       raw_->read(remaining_bytes, reinterpret_cast<uint8_t *>(out) +
                                                                   pre_buffer_copy_bytes));
                raw_read_total_ += bytes_read;
                RewindBuffer();
                return pre_buffer_copy_bytes + bytes_read;
            } else {
                // 2.2. If read is smaller than buffer size, fill buffer and copy from buffer.
                TURBO_RETURN_NOT_OK(DoBuffer());
                int64_t bytes_copy_after_buffer = std::min(bytes_buffered_, remaining_bytes);
                memcpy(reinterpret_cast<uint8_t *>(out) + pre_buffer_copy_bytes,
                       buffer_data_ + buffer_pos_, bytes_copy_after_buffer);
                ConsumeBuffer(bytes_copy_after_buffer);
                return pre_buffer_copy_bytes + bytes_copy_after_buffer;
            }
        }

        turbo::Result<std::shared_ptr<Buffer>> read(int64_t nbytes) {
            TURBO_MOVE_OR_RAISE(auto buffer, allocate_resizable_buffer(nbytes, pool_));

            TURBO_MOVE_OR_RAISE(int64_t bytes_read, read(nbytes, buffer->mutable_data()));

            if (bytes_read < nbytes) {
                // Change size but do not reallocate internal capacity
                TURBO_RETURN_NOT_OK(buffer->resize(bytes_read, false /* shrink_to_fit */));
                buffer->zero_padding();
            }
            // R build with openSUSE155 requires an explicit shared_ptr construction
            return std::shared_ptr<Buffer>(std::move(buffer));
        }

        // For providing access to the raw file handles
        std::shared_ptr<InputStream> raw() const { return raw_; }

    private:
        std::shared_ptr<InputStream> raw_;
        int64_t raw_read_total_;
        // a bound on the maximum number of bytes to read from the raw input stream.
        // The default -1 indicates that it is unbounded
        int64_t raw_read_bound_;

        // Number of remaining bytes in the buffer, to be reduced on each read from
        // the buffer
        int64_t bytes_buffered_;
    };

    BufferedInputStream::BufferedInputStream(std::shared_ptr<InputStream> raw,
                                             MemoryPool *pool,
                                             int64_t raw_total_bytes_bound) {
        impl_ = std::make_unique<Impl>(std::move(raw), pool, raw_total_bytes_bound);
    }

    BufferedInputStream::~BufferedInputStream() { internal::CloseFromDestructor(this); }

    turbo::Result<std::shared_ptr<BufferedInputStream>> BufferedInputStream::create(
            int64_t buffer_size, MemoryPool *pool, std::shared_ptr<InputStream> raw,
            int64_t raw_total_bytes_bound) {
        auto result = std::shared_ptr<BufferedInputStream>(
                new BufferedInputStream(std::move(raw), pool, raw_total_bytes_bound));
        TURBO_RETURN_NOT_OK(result->set_buffer_size(buffer_size));
        return result;
    }

    turbo::Status BufferedInputStream::DoClose() { return impl_->close(); }

    turbo::Status BufferedInputStream::DoAbort() { return impl_->abort(); }

    bool BufferedInputStream::closed() const { return impl_->closed(); }

    std::shared_ptr<InputStream> BufferedInputStream::detach() { return impl_->detach(); }

    std::shared_ptr<InputStream> BufferedInputStream::raw() const { return impl_->raw(); }

    turbo::Result<int64_t> BufferedInputStream::DoTell() const { return impl_->tell(); }

    turbo::Result<std::string_view> BufferedInputStream::DoPeek(int64_t nbytes) {
        return impl_->peek(nbytes);
    }

    turbo::Status BufferedInputStream::set_buffer_size(int64_t new_buffer_size) {
        return impl_->set_buffer_size(new_buffer_size);
    }

    int64_t BufferedInputStream::bytes_buffered() const { return impl_->bytes_buffered(); }

    int64_t BufferedInputStream::buffer_size() const { return impl_->buffer_size(); }

    turbo::Result<int64_t> BufferedInputStream::DoRead(int64_t nbytes, void *out) {
        return impl_->read(nbytes, out);
    }

    turbo::Result<std::shared_ptr<Buffer>> BufferedInputStream::DoRead(int64_t nbytes) {
        return impl_->read(nbytes);
    }

    turbo::Result<std::shared_ptr<const KeyValueMetadata>> BufferedInputStream::read_metadata() {
        return impl_->raw()->read_metadata();
    }

    Future<std::shared_ptr<const KeyValueMetadata>> BufferedInputStream::read_metadata_async(
            const IOContext &io_context) {
        return impl_->raw()->read_metadata_async(io_context);
    }

}  // namespace nebula::io

